"""
数据库连接器工具
支持多种数据库类型的统一连接接口
"""

import asyncio
import time
from typing import Dict, Any, Optional, Tuple, List
from contextlib import asynccontextmanager
import logging

from app.models.datasource import DatabaseType
from app.schemas.datasource import ConnectionTestRequest, ConnectionTestResponse
from app.utils.kafka_client import KafkaClient

logger = logging.getLogger(__name__)


class DatabaseConnectionError(Exception):
    """数据库连接异常"""
    pass


class DatabaseConnector:
    """数据库连接器基类"""
    
    def __init__(self, db_type: DatabaseType, host: str, port: int, 
                 database: str, username: str, password: str,
                 connection_params: Optional[Dict[str, Any]] = None,
                 timeout: int = 30):
        self.db_type = db_type
        self.host = host
        self.port = port
        self.database = database
        self.username = username
        self.password = password
        self.connection_params = connection_params or {}
        self.timeout = timeout
    
    async def test_connection(self) -> ConnectionTestResponse:
        """测试数据库连接"""
        start_time = time.time()
        
        try:
            # 根据数据库类型选择连接方法
            if self.db_type == DatabaseType.MYSQL:
                result = await self._test_mysql_connection()
            elif self.db_type == DatabaseType.POSTGRESQL:
                result = await self._test_postgresql_connection()
            elif self.db_type == DatabaseType.ORACLE:
                result = await self._test_oracle_connection()
            elif self.db_type == DatabaseType.SQLSERVER:
                result = await self._test_sqlserver_connection()
            elif self.db_type == DatabaseType.DORIS:
                result = await self._test_doris_connection()
            elif self.db_type == DatabaseType.KAFKA:
                result = await self._test_kafka_connection()
            else:
                raise DatabaseConnectionError(f"不支持的数据库类型: {self.db_type}")
            
            response_time = int((time.time() - start_time) * 1000)
            
            return ConnectionTestResponse(
                success=True,
                message="连接测试成功",
                response_time=response_time,
                database_info=result
            )
            
        except Exception as e:
            response_time = int((time.time() - start_time) * 1000)
            error_message = str(e)
            
            logger.error(f"数据库连接测试失败 [{self.db_type}://{self.host}:{self.port}]: {error_message}")
            
            return ConnectionTestResponse(
                success=False,
                message=f"连接失败: {error_message}",
                response_time=response_time,
                error_details={
                    "error_type": type(e).__name__,
                    "error_message": error_message,
                    "database_type": self.db_type,
                    "host": self.host,
                    "port": self.port
                }
            )
    
    async def _test_mysql_connection(self) -> Dict[str, Any]:
        """测试MySQL连接"""
        try:
            import aiomysql
            
            connection = await asyncio.wait_for(
                aiomysql.connect(
                    host=self.host,
                    port=self.port,
                    user=self.username,
                    password=self.password,
                    db=self.database,
                    **self.connection_params
                ),
                timeout=self.timeout
            )
            
            async with connection:
                cursor = await connection.cursor()
                await cursor.execute("SELECT VERSION() as version")
                result = await cursor.fetchone()
                version = result[0] if result else "Unknown"
                
                await cursor.execute("SELECT DATABASE() as current_db")
                result = await cursor.fetchone()
                current_db = result[0] if result else "Unknown"
                
                return {
                    "database_type": "MySQL",
                    "version": version,
                    "current_database": current_db,
                    "server_info": connection.get_server_info() if hasattr(connection, 'get_server_info') else None
                }
                
        except ImportError:
            raise DatabaseConnectionError("MySQL驱动未安装，请安装aiomysql")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"连接超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"MySQL连接失败: {str(e)}")
    
    async def _test_postgresql_connection(self) -> Dict[str, Any]:
        """测试PostgreSQL连接"""
        try:
            import asyncpg
            
            connection = await asyncio.wait_for(
                asyncpg.connect(
                    host=self.host,
                    port=self.port,
                    user=self.username,
                    password=self.password,
                    database=self.database,
                    **self.connection_params
                ),
                timeout=self.timeout
            )
            
            async with connection:
                version = await connection.fetchval("SELECT version()")
                current_db = await connection.fetchval("SELECT current_database()")
                
                return {
                    "database_type": "PostgreSQL",
                    "version": version,
                    "current_database": current_db,
                    "server_version": connection.get_server_version()
                }
                
        except ImportError:
            raise DatabaseConnectionError("PostgreSQL驱动未安装，请安装asyncpg")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"连接超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"PostgreSQL连接失败: {str(e)}")
    
    async def _test_oracle_connection(self) -> Dict[str, Any]:
        """测试Oracle连接"""
        try:
            import cx_Oracle
            
            # Oracle连接字符串
            dsn = cx_Oracle.makedsn(self.host, self.port, service_name=self.database)
            
            # 由于cx_Oracle不支持异步，使用线程池执行
            loop = asyncio.get_event_loop()
            
            def _sync_connect():
                connection = cx_Oracle.connect(
                    user=self.username,
                    password=self.password,
                    dsn=dsn,
                    **self.connection_params
                )
                
                with connection:
                    cursor = connection.cursor()
                    cursor.execute("SELECT * FROM v$version WHERE banner LIKE 'Oracle%'")
                    result = cursor.fetchone()
                    version = result[0] if result else "Unknown"
                    
                    cursor.execute("SELECT sys_context('USERENV', 'DB_NAME') FROM dual")
                    result = cursor.fetchone()
                    current_db = result[0] if result else "Unknown"
                    
                    return {
                        "database_type": "Oracle",
                        "version": version,
                        "current_database": current_db,
                        "client_version": cx_Oracle.clientversion()
                    }
            
            return await asyncio.wait_for(
                loop.run_in_executor(None, _sync_connect),
                timeout=self.timeout
            )
            
        except ImportError:
            raise DatabaseConnectionError("Oracle驱动未安装，请安装cx_Oracle")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"连接超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"Oracle连接失败: {str(e)}")
    
    async def _test_sqlserver_connection(self) -> Dict[str, Any]:
        """测试SQL Server连接"""
        try:
            import pyodbc
            
            # SQL Server连接字符串
            connection_string = (
                f"DRIVER={{ODBC Driver 17 for SQL Server}};"
                f"SERVER={self.host},{self.port};"
                f"DATABASE={self.database};"
                f"UID={self.username};"
                f"PWD={self.password};"
            )
            
            # 添加额外参数
            for key, value in self.connection_params.items():
                connection_string += f"{key}={value};"
            
            # 由于pyodbc不支持异步，使用线程池执行
            loop = asyncio.get_event_loop()
            
            def _sync_connect():
                connection = pyodbc.connect(connection_string, timeout=self.timeout)
                
                with connection:
                    cursor = connection.cursor()
                    cursor.execute("SELECT @@VERSION as version")
                    result = cursor.fetchone()
                    version = result[0] if result else "Unknown"
                    
                    cursor.execute("SELECT DB_NAME() as current_db")
                    result = cursor.fetchone()
                    current_db = result[0] if result else "Unknown"
                    
                    return {
                        "database_type": "SQL Server",
                        "version": version,
                        "current_database": current_db
                    }
            
            return await asyncio.wait_for(
                loop.run_in_executor(None, _sync_connect),
                timeout=self.timeout
            )
            
        except ImportError:
            raise DatabaseConnectionError("SQL Server驱动未安装，请安装pyodbc")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"连接超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"SQL Server连接失败: {str(e)}")
    
    async def _test_doris_connection(self) -> Dict[str, Any]:
        """测试Doris连接（使用MySQL协议）"""
        try:
            import aiomysql
            
            connection = await asyncio.wait_for(
                aiomysql.connect(
                    host=self.host,
                    port=self.port,
                    user=self.username,
                    password=self.password,
                    db=self.database,
                    **self.connection_params
                ),
                timeout=self.timeout
            )
            
            async with connection:
                cursor = await connection.cursor()
                await cursor.execute("SELECT @@version_comment as version")
                result = await cursor.fetchone()
                version = result[0] if result else "Unknown"

                await cursor.execute("SELECT DATABASE() as current_db")
                result = await cursor.fetchone()
                current_db = result[0] if result else "Unknown"
                
                return {
                    "database_type": "Apache Doris",
                    "version": version,
                    "current_database": current_db
                }
                
        except ImportError:
            raise DatabaseConnectionError("Doris驱动未安装，请安装aiomysql")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"连接超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"Doris连接失败: {str(e)}")

    async def _test_kafka_connection(self) -> Dict[str, Any]:
        """测试Kafka连接"""
        try:
            # 构建Kafka服务器地址
            bootstrap_servers = f"{self.host}:{self.port}"

            # 创建Kafka客户端
            kafka_client = KafkaClient(bootstrap_servers, **self.connection_params)

            # 测试连接
            test_result = await kafka_client.test_connection(timeout=self.timeout)

            if test_result['success']:
                return {
                    "message_broker": "Apache Kafka",
                    "bootstrap_servers": bootstrap_servers,
                    "cluster_info": test_result.get('cluster_info', {})
                }
            else:
                raise DatabaseConnectionError(test_result['message'])

        except Exception as e:
            raise DatabaseConnectionError(f"Kafka连接失败: {str(e)}")

    async def execute_query(self, sql: str) -> List[Dict[str, Any]]:
        """
        执行SQL查询

        Args:
            sql: SQL查询语句

        Returns:
            查询结果列表
        """
        try:
            # 根据数据库类型选择执行方法
            if self.db_type == DatabaseType.MYSQL:
                return await self._execute_mysql_query(sql)
            elif self.db_type == DatabaseType.POSTGRESQL:
                return await self._execute_postgresql_query(sql)
            elif self.db_type == DatabaseType.ORACLE:
                return await self._execute_oracle_query(sql)
            elif self.db_type == DatabaseType.SQLSERVER:
                return await self._execute_sqlserver_query(sql)
            elif self.db_type == DatabaseType.DORIS:
                return await self._execute_doris_query(sql)
            else:
                raise DatabaseConnectionError(f"不支持的数据库类型: {self.db_type}")

        except Exception as e:
            logger.error(f"执行SQL查询失败 [{self.db_type}]: {sql[:100]}... - {str(e)}")
            raise DatabaseConnectionError(f"SQL查询执行失败: {str(e)}")

    async def _execute_mysql_query(self, sql: str) -> List[Dict[str, Any]]:
        """执行MySQL查询"""
        try:
            import aiomysql

            connection = await asyncio.wait_for(
                aiomysql.connect(
                    host=self.host,
                    port=self.port,
                    user=self.username,
                    password=self.password,
                    db=self.database,
                    **self.connection_params
                ),
                timeout=self.timeout
            )

            async with connection:
                cursor = await connection.cursor(aiomysql.DictCursor)
                # 添加查询超时，避免长时间卡住
                await asyncio.wait_for(cursor.execute(sql), timeout=60)  # 60秒超时
                result = await asyncio.wait_for(cursor.fetchall(), timeout=60)  # 60秒超时
                return [dict(row) for row in result] if result else []

        except ImportError:
            raise DatabaseConnectionError("MySQL驱动未安装，请安装aiomysql")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"查询超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"MySQL查询失败: {str(e)}")

    async def _execute_postgresql_query(self, sql: str) -> List[Dict[str, Any]]:
        """执行PostgreSQL查询"""
        try:
            import asyncpg

            connection = await asyncio.wait_for(
                asyncpg.connect(
                    host=self.host,
                    port=self.port,
                    user=self.username,
                    password=self.password,
                    database=self.database,
                    **self.connection_params
                ),
                timeout=self.timeout
            )

            async with connection:
                result = await connection.fetch(sql)
                return [dict(row) for row in result] if result else []

        except ImportError:
            raise DatabaseConnectionError("PostgreSQL驱动未安装，请安装asyncpg")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"查询超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"PostgreSQL查询失败: {str(e)}")

    async def _execute_oracle_query(self, sql: str) -> List[Dict[str, Any]]:
        """执行Oracle查询"""
        try:
            import cx_Oracle

            # Oracle连接字符串
            dsn = cx_Oracle.makedsn(self.host, self.port, service_name=self.database)

            # 由于cx_Oracle不支持异步，使用线程池执行
            loop = asyncio.get_event_loop()

            def _sync_query():
                connection = cx_Oracle.connect(
                    user=self.username,
                    password=self.password,
                    dsn=dsn,
                    **self.connection_params
                )

                with connection:
                    cursor = connection.cursor()
                    cursor.execute(sql)

                    # 获取列名
                    columns = [desc[0] for desc in cursor.description] if cursor.description else []

                    # 获取数据
                    rows = cursor.fetchall()

                    # 转换为字典列表
                    result = []
                    for row in rows:
                        result.append(dict(zip(columns, row)))

                    return result

            return await asyncio.wait_for(
                loop.run_in_executor(None, _sync_query),
                timeout=self.timeout
            )

        except ImportError:
            raise DatabaseConnectionError("Oracle驱动未安装，请安装cx_Oracle")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"查询超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"Oracle查询失败: {str(e)}")

    async def _execute_sqlserver_query(self, sql: str) -> List[Dict[str, Any]]:
        """执行SQL Server查询"""
        try:
            import pyodbc

            # SQL Server连接字符串
            connection_string = (
                f"DRIVER={{ODBC Driver 17 for SQL Server}};"
                f"SERVER={self.host},{self.port};"
                f"DATABASE={self.database};"
                f"UID={self.username};"
                f"PWD={self.password};"
            )

            # 添加额外参数
            for key, value in self.connection_params.items():
                connection_string += f"{key}={value};"

            # 由于pyodbc不支持异步，使用线程池执行
            loop = asyncio.get_event_loop()

            def _sync_query():
                connection = pyodbc.connect(connection_string, timeout=self.timeout)

                with connection:
                    cursor = connection.cursor()
                    cursor.execute(sql)

                    # 获取列名
                    columns = [column[0] for column in cursor.description] if cursor.description else []

                    # 获取数据
                    rows = cursor.fetchall()

                    # 转换为字典列表
                    result = []
                    for row in rows:
                        result.append(dict(zip(columns, row)))

                    return result

            return await asyncio.wait_for(
                loop.run_in_executor(None, _sync_query),
                timeout=self.timeout
            )

        except ImportError:
            raise DatabaseConnectionError("SQL Server驱动未安装，请安装pyodbc")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"查询超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"SQL Server查询失败: {str(e)}")

    async def _execute_doris_query(self, sql: str) -> List[Dict[str, Any]]:
        """执行Doris查询（使用MySQL协议）"""
        try:
            import aiomysql

            connection = await asyncio.wait_for(
                aiomysql.connect(
                    host=self.host,
                    port=self.port,
                    user=self.username,
                    password=self.password,
                    db=self.database,
                    **self.connection_params
                ),
                timeout=self.timeout
            )

            async with connection:
                cursor = await connection.cursor(aiomysql.DictCursor)
                await cursor.execute(sql)
                result = await cursor.fetchall()
                return [dict(row) for row in result] if result else []

        except ImportError:
            raise DatabaseConnectionError("Doris驱动未安装，请安装aiomysql")
        except asyncio.TimeoutError:
            raise DatabaseConnectionError(f"查询超时（{self.timeout}秒）")
        except Exception as e:
            raise DatabaseConnectionError(f"Doris查询失败: {str(e)}")


async def test_database_connection(request: ConnectionTestRequest) -> ConnectionTestResponse:
    """测试数据库连接的便捷函数"""
    try:
        connector = DatabaseConnector(
            db_type=request.type,
            host=request.host,
            port=request.port,
            database=request.database_name,
            username=request.username,
            password=request.password,
            connection_params=request.connection_params,
            timeout=request.connection_timeout
        )
        
        return await connector.test_connection()
        
    except Exception as e:
        logger.error(f"数据库连接测试失败: {str(e)}")
        return ConnectionTestResponse(
            success=False,
            message=f"连接测试失败: {str(e)}",
            error_details={
                "error_type": type(e).__name__,
                "error_message": str(e)
            }
        )


def get_default_port(db_type: DatabaseType) -> int:
    """获取数据库类型的默认端口"""
    default_ports = {
        DatabaseType.MYSQL: 3306,
        DatabaseType.POSTGRESQL: 5432,
        DatabaseType.ORACLE: 1521,
        DatabaseType.SQLSERVER: 1433,
        DatabaseType.DORIS: 9030,  # MySQL协议端口，用于查询
    }
    return default_ports.get(db_type, 3306)


def get_doris_http_port(mysql_port: int) -> int:
    """根据Doris MySQL端口获取对应的HTTP端口"""
    # 常见的端口映射关系
    port_mapping = {
        9030: 18030,  # 用户环境的端口映射
        9031: 18031,  # 可能的其他实例
        9032: 18032,  # 可能的其他实例
    }

    # 如果有明确的映射关系，使用映射
    if mysql_port in port_mapping:
        return port_mapping[mysql_port]

    # 如果已经是HTTP端口，直接返回
    if mysql_port in [8030, 8031, 8032, 18030, 18031, 18032]:
        return mysql_port

    # 默认使用18030
    return 18030


def get_connection_string_template(db_type: DatabaseType) -> str:
    """获取数据库连接字符串模板"""
    templates = {
        DatabaseType.MYSQL: "mysql://{username}:{password}@{host}:{port}/{database}",
        DatabaseType.POSTGRESQL: "postgresql://{username}:{password}@{host}:{port}/{database}",
        DatabaseType.ORACLE: "oracle://{username}:{password}@{host}:{port}/{database}",
        DatabaseType.SQLSERVER: "mssql://{username}:{password}@{host}:{port}/{database}",
        DatabaseType.DORIS: "mysql://{username}:{password}@{host}:{port}/{database}",
    }
    return templates.get(db_type, "unknown://{username}:{password}@{host}:{port}/{database}")


# 测试函数
async def test_all_database_types():
    """测试所有数据库类型的连接（仅用于开发测试）"""
    test_configs = [
        {
            "type": DatabaseType.MYSQL,
            "host": "localhost",
            "port": 3306,
            "database_name": "test",
            "username": "root",
            "password": "password"
        }
        # 可以添加更多测试配置
    ]
    
    for config in test_configs:
        try:
            request = ConnectionTestRequest(**config)
            result = await test_database_connection(request)
            print(f"✅ {config['type']} 连接测试: {result.message}")
        except Exception as e:
            print(f"❌ {config['type']} 连接测试失败: {str(e)}")


if __name__ == "__main__":
    # 运行测试
    asyncio.run(test_all_database_types())
